#!/bin/bash

# Check if sq_task_pids file is provided


# 定义排序函数
sort_pids_by_dev() {
    local input_file="$1"
    local output_file="$2"
    
    # 参数检查
    if [[ -z "$input_file" || -z "$output_file" ]]; then
        echo "Usage: sort_pids_by_dev <input_file> <output_file>"
        return 1
    fi
    
    # 检查输入文件是否存在
    if [[ ! -f "$input_file" ]]; then
        echo "Error: Input file '$input_file' not found"
        return 1
    fi
    
    # 处理文件内容：
    # 1. 使用awk提取设备号和整行内容
    # 2. 按设备号排序
    # 3. 提取第一列(PID)
    awk '{
        # 提取设备号
        match($2, /dev([0-9]+)_sq_task/, arr)
        if (length(arr[1]) > 0) {
            dev_num = int(arr[1])
            print dev_num, $0
        }
    }' "$input_file" | \
    sort -n -k1,1 | \
    awk '{print $2}' > "$output_file"
    
    echo "排序完成，结果已写入 $output_file"
}


ps -ef | grep _sq_task | awk '{if (NR>1) print prev; prev=$0;}' | awk '{print $2, $8}' > tmp.txt

sort_pids_by_dev "tmp.txt" "tmp2.txt"


SQ_TASK_PIDS_FILE=tmp2.txt

# Check if the file exists
if [ ! -f "$SQ_TASK_PIDS_FILE" ]; then
    echo "Error: File $SQ_TASK_PIDS_FILE not found!"
    exit 1
fi

decimal_to_320bit_hex_mask() {
    local decimal_input="$1"

    # --- 参数验证 ---
    if [ $# -ne 1 ]; then
        echo "错误 (函数): 请提供一个参数。" >&2
        return 1
    fi

    if ! [[ "$decimal_input" =~ ^[0-9]+$ ]]; then
        echo "错误 (函数): '$decimal_input' 不是一个有效的非负整数。" >&2
        return 1
    fi

    if [ "$decimal_input" -lt 0 ] || [ "$decimal_input" -gt 319 ]; then
        echo "错误 (函数): 输入值 $decimal_input 超出范围 (0-319)。" >&2
        return 1
    fi

    # --- 核心计算逻辑 ---
    local BITS_PER_GROUP=32
    local NUM_GROUPS=10  # 320 bits / 32 bits per group = 10 groups

    # 计算组索引 (从左到右 0-9) 和组内位索引 (0-31)
    local group_index=$((decimal_input / BITS_PER_GROUP))
    local bit_position_in_group=$((decimal_input % BITS_PER_GROUP))

    # 计算该组的十进制值 (2^bit_position_in_group)
    local group_value_decimal=$((1 << bit_position_in_group))

    # 将十进制值转换为8位十六进制 (大写)
    local group_value_hex
    group_value_hex=$(printf "%08X" "$group_value_decimal")

    # --- 构建并输出掩码 ---
    local output_mask=""
    local i

    for (( i=0; i<NUM_GROUPS; i++ )); do
        if [ $i -eq $((NUM_GROUPS - 1 - group_index)) ]; then
            # 注意组的顺序：最高位组在最左边，所以我们要在正确的位置放置非零组
            # 组索引是从0开始的，而输出顺序是最高位组在前
            # 如果bit在第group_index组（从0开始计算），则在输出时应为第(NUM_GROUPS-1-group_index)组
            output_mask="${output_mask}${group_value_hex}"
        else
            output_mask="${output_mask}00000000"
        fi

        # 添加逗号分隔符（除了最后一个）
        if [ $i -lt $((NUM_GROUPS - 1)) ]; then
            output_mask="${output_mask},"
        fi
    done

    echo "$output_mask"
    return 0
}

# 根据cluser 第二个获取CPU核心列表
get_cluster_second_cpu() {
    local numa_id=$1
    local cluster_id=$2
    
    # 参数验证
    if [[ -z "$numa_id" || -z "$cluster_id" ]]; then
        echo "错误: 需要提供 numa_id 和 cluster_id 参数" >&2
        return 1
    fi
    
    # 验证参数是否为数字
    if ! [[ "$numa_id" =~ ^[0-9]+$ && "$cluster_id" =~ ^[0-9]+$ ]]; then
        echo "错误: 参数必须为非负整数" >&2
        return 1
    fi
    
    local second_cpu=$((numa_id * 40 + cluster_id * 4 + 1))

    echo "$second_cpu"
}

# 根据cluser 第三个获取CPU核心列表
get_cluster_third_cpu() {
    local numa_id=$1
    local cluster_id=$2
    
    # 参数验证
    if [[ -z "$numa_id" || -z "$cluster_id" ]]; then
        echo "错误: 需要提供 numa_id 和 cluster_id 参数" >&2
        return 1
    fi
    
    # 验证参数是否为数字
    if ! [[ "$numa_id" =~ ^[0-9]+$ && "$cluster_id" =~ ^[0-9]+$ ]]; then
        echo "错误: 参数必须为非负整数" >&2
        return 1
    fi
    
    local third_cpu=$((numa_id * 40 + cluster_id * 4 + 2))

    echo "$third_cpu"
}


# Initialize cpu-no with the base value

task_id=0

# Process each PID in the sq_task_pids file
while IFS= read -r sq_task_pid; do
    # Skip empty lines
    if [ -n "$sq_task_pid" ]; then
        local numa_id=$((task_id / 8) + 6) # NUMA 6-7
        local cluster_id=$((task_id % 8)) # CLUSTER 0-7
        local cpu_no=$(get_cluster_second_cpu "numa_id" "cluster_id")

        echo "Executing: taskset -pc $cpu_no $sq_task_pid"
        taskset -pc $cpu_no $sq_task_pid

        task_id=$((task_id + 1))
    fi
done < "$SQ_TASK_PIDS_FILE"

# Get all interrupt numbers for trs-mb
echo "Getting interrupt numbers for trs-mb..."
IRQ_NUMBERS=$(cat /proc/interrupts | grep trs-mb | awk -F ":" '{print $1}' | tr -d ' ')

IRQ_CPU=6
SQ_CPU=12

task_id=0
# Process each interrupt number
for irq_no in $IRQ_NUMBERS; do
    local numa_id=$((task_id / 8) + 6) # NUMA 6-7
    local cluster_id=$((task_id % 8)) # CLUSTER 0-7
    local cpu_no=$(get_cluster_third_cpu "numa_id" "cluster_id")

    echo "Executing: echo $cpu_no > /proc/irq/$irq_no/smp_affinity_list"
    echo $cpu_no > /proc/irq/$irq_no/smp_affinity_list

    task_id=$((task_id + 1))
done

task_id=0
# Process dev-no from 1 to 16
for dev_no in $(seq 0 15); do
    # Execute for _sq_send_wq
    local numa_id=$((task_id / 8) + 6) # NUMA 6-7
    local cluster_id=$((task_id % 8)) # CLUSTER 0-7
    local cpu_no=$(get_cluster_third_cpu "numa_id" "cluster_id")

    local cpu_mask=$(decimal_to_320bit_hex_mask $cpu_no)

    echo "Executing: echo $cpu_mask > /sys/devices/virtual/workqueue/dev${dev_no}_sq_send_wq/cpumask"
    echo $cpu_mask > /sys/devices/virtual/workqueue/dev${dev_no}_sq_send_wq/cpumask
    
    # Execute for _sq_fair_wq
    echo "Executing: echo $cpu_mask > /sys/devices/virtual/workqueue/dev${dev_no}_sq_fair_wq/cpumask"
    echo $cpu_mask > /sys/devices/virtual/workqueue/dev${dev_no}_sq_fair_wq/cpumask
    
    task_id=$((task_id + 1))
done

echo "Script execution completed."




